﻿#include"fbank.h"
#include<math.h>
#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<assert.h>
//#include"fftr.h"
#ifndef M_PI
#define M_PI       3.14159265358979323846
#endif

void hanning(float* win, int16_t len) {
	for (int16_t i = 0; i < len; ++i) {
		win[i] = 0.5 - 0.5 * cosf(2 * M_PI * i / (len - 1));
	}
}
void sine(float* win, int16_t len) {
	for (int16_t i = 0; i < len; ++i) {
		win[i] = sinf(M_PI*i/(len-1));
	}
}
void hamming(float* win, int16_t len) {
	for (int16_t i = 0; i < len; ++i) {
		win[i] = 0.54 - 0.46 * cosf(2 * M_PI * i / (len - 1));
	}
}
void povey(float* win, int16_t len) {
	for (int16_t i = 0; i < len; ++i) {
		win[i] = powf(0.5 - 0.5 * cosf(2 * M_PI * i / (len - 1)), 0.85);
	}
}
void rectangular(float* win, int16_t len) {
	for (int16_t i = 0; i < len; ++i) {
		win[i] = 1.0f;
	}
}
void blackman(float* win, int16_t len,float blackman_coeff) {
	for (int16_t i = 0; i < len; ++i) {
		win[i] = blackman_coeff - 0.5 * cosf(2 * M_PI * i / (len - 1)) + (0.5 - blackman_coeff) * cosf(4 * M_PI * i / (len - 1));
	}
}
void getwin(float* win, int16_t len, wintype type, float blackman_coeff) {
	switch (type){
	case HANNING:
		hanning(win, len);
		break;
	case SINE:
		sine(win, len);
		break;
	case HAMMING:
		hamming(win, len);
		break;
	case POVEY:
		povey(win, len);
		break;
	case RECTANGULAR:
		rectangular(win, len);
		break;
	case BLACKMAN:
		blackman(win, len, blackman_coeff);
	default:
		break;
	}
}

int16_t round_to_power_of_two(int16_t x) {
	int16_t y = 1;
	for (int i = 0; i < 15; ++i) {
		if (y >= x) {
			return y;
		}
		y <<= 1;
	}
	return -1;
}

float mel_scale_scalar(float freq) {
	return 1127.0f * logf(1.0 + freq / 700);
}
float inverse_mel_scale(float mel_freq) {
	return 700.0 * (expf((mel_freq / 1127) - 1.0));
}

float get_log_energy(float* input, float epsilon, float energy_floor,int16_t len) {
	float sum = 0;
	for (int i = 0; i < len; ++i) {
		sum += input[i] * input[i];
	}
	sum = sum > epsilon ? sum : epsilon;
	sum = sum > energy_floor ? sum : energy_floor;
	return logf(sum);
}
static void hadamard(float* a, float* b, float* out, int16_t len) {
	for (int i = 0; i < len; ++i) {
		out[i] = a[i] * b[i];
	}
}
static void vec_dot(float* a, float* b, int16_t len,float* out) {
	for (int i = 0; i < len; ++i) {
		*out += a[i] * b[i];
	}
}


// vtln技术暂时不支持
static float** get_mel_banks(int num_bins, int16_t window_length_padded, float sample_freq,
	float low_freq, float high_freq, int16_t* bin_start, int16_t* bin_len) {
	assert(num_bins > 3);
	assert(window_length_padded % 2 == 0);
	int16_t num_fft_bins = 0.5* window_length_padded;
	float nyquist = 0.5 * sample_freq;
	if (high_freq <= 0.0) {
		high_freq += nyquist;
	}
	assert((0.0 <= low_freq < nyquist) && (0.0 < high_freq <= nyquist) && (low_freq < high_freq));
	float fft_bin_width = sample_freq / window_length_padded;
	float mel_low_freq = mel_scale_scalar(low_freq);
	float mel_high_freq = mel_scale_scalar(high_freq);
	float mel_freq_delta = (mel_high_freq - mel_low_freq) / (num_bins + 1);
	//if (vtln_high < 0.0) {
		//vtln_high += nyquist;
	//}
	//assert(vtln_warp_factor == 1.0 || ((low_freq < vtln_low < high_freq) && (0.0 < vtln_high < high_freq) && (vtln_low < vtln_high)));
	
	float* left_mel = (float*)malloc(sizeof(float) * num_bins);
	float* center_mel = (float*)malloc(sizeof(float) * num_bins);
	float* right_mel = (float*)malloc(sizeof(float) * num_bins);
	float* mel= (float*)malloc(sizeof(float) * num_fft_bins);
	int i = 0;
	int j = 0;
	for (i = 0; i < num_bins; ++i) {
		left_mel[i] = mel_low_freq + i * mel_freq_delta;
		center_mel[i] = mel_low_freq + (i + 1.0) * mel_freq_delta;
		right_mel[i] = mel_low_freq + (i + 2.0) * mel_freq_delta;
	}
	for (i = 0; i < num_fft_bins; ++i) {
		mel[i] = mel_scale_scalar(fft_bin_width * i);
	}
	int coefcnt = 0;
	for (i = 0; i < num_bins; ++i) {
		bin_len[i] = 0;
		bin_start[i] = -1;
		for (j = 0; j < num_fft_bins; ++j) {
			if (mel[j] - left_mel[i] > 0 && right_mel[i] - mel[j] > 0) {
				bin_len[i]++;
				coefcnt++;
				if (bin_start[i] < 0) {
					bin_start[i] = j;	
			}
			}
		}
	}
	float* coefdata = (float*)malloc(sizeof(float) * coefcnt);
	float** coef = (float**)malloc(sizeof(float*) * num_bins);
	float* relu = coefdata;
	float l = 0, r = 0;
	for (i = 0; i < num_bins; ++i) {
		coef[i] = relu;
		relu += bin_len[i];
		for (j = bin_start[i]; j < bin_start[i]+ bin_len[i]; ++j) {
			l = (mel[j] - left_mel[i]) / (center_mel[i] - left_mel[i]);
			r = (right_mel[i] - mel[j]) / (right_mel[i] - center_mel[i]);
			*coefdata = l > r ? r : l;
			coefdata++;
		}
	}
	free(left_mel);
	free(center_mel);
	free(right_mel);
	free(mel);
	return coef;
}

fbank* FbankInit(uint16_t num_mel_bins, int16_t sample_frequency) {
	fbank* inst = (fbank*)malloc(sizeof(fbank));
	if (inst) {
		inst->sample_frequency = sample_frequency;
		inst->num_mel_bins = num_mel_bins;
		inst->blackman_coeff = 0.42;
		inst->energy_floor = 1.0;
		inst->frame_length = 25;
		inst->frame_shift = 10;
		inst->high_freq = 0.0;
		inst->low_freq = 20.0;
		inst->htk_compat = 0;
		inst->preemphasis_coefficient = 0.97;
		inst->raw_energy = 1;
		inst->remove_dc_offset = 1;
		inst->round_to_power_of_two = 1;
		inst->subtract_mean = 0;
		inst->use_energy = 0;
		inst->use_log_fbank = 1;
		inst->use_power = 1;
		inst->window_type = POVEY;

		inst->bin_start = (int16_t*)malloc(sizeof(int16_t) * num_mel_bins);
		inst->bin_len = (int16_t*)malloc(sizeof(int16_t) * num_mel_bins);
		inst->window_size = inst->frame_length * inst->sample_frequency / 1000;
		inst->window = (float*)malloc(sizeof(float) * inst->window_size);
	}
	return inst;
}

void FbankRegister(fbank* inst) {
	getwin(inst->window, inst->window_size, inst->window_type, inst->blackman_coeff);
	if (inst->round_to_power_of_two) {
		inst->padded_window_size = round_to_power_of_two(inst->window_size);
	}
	else {
		inst->padded_window_size = inst->window_size;
	}
	inst->table_td = (float*)malloc(sizeof(float) * inst->padded_window_size);
	inst->fft = fftInit(inst->padded_window_size);
	inst->table_fd = (fft_cpx*)malloc(sizeof(fft_cpx) * (inst->padded_window_size / 2 + 1));
	
	inst->melbins = get_mel_banks(inst->num_mel_bins,
		inst->padded_window_size,
		inst->sample_frequency,
		inst->low_freq,
		inst->high_freq,
		inst->bin_start,
		inst->bin_len
	);
	if (inst->use_energy) {
		inst->fbank_bins = inst->num_mel_bins + 1;
	}
	else {
		inst->fbank_bins = inst->num_mel_bins;
	}
}

void FbankProcess(fbank* inst, float* frame_in, float* fbank_out) {
	memset(inst->table_td, 0, sizeof(float) * inst->padded_window_size);
	memset(fbank_out, 0, sizeof(float) * inst->fbank_bins);
	memcpy(inst->table_td, frame_in, sizeof(float) * inst->window_size);
	if (inst->remove_dc_offset) {
		float mean = 0;
		for (int i = 0; i < inst->window_size; ++i) {
			mean += inst->table_td[i];
		}
		mean /= inst->window_size;
		for (int i = 0; i < inst->window_size; ++i) {
			inst->table_td[i] -= mean;
		}
	}
	float signal_log_energy = 0;
	float eps = 1.1920928955078125e-07f;
	if (inst->raw_energy) {
		signal_log_energy = get_log_energy(inst->table_td, eps, inst->energy_floor, inst->window_size);
	}
	if (inst->preemphasis_coefficient != 0.0f) {
		for (int i = inst->window_size - 1; i >= 1; --i) {
			inst->table_td[i] = inst->table_td[i] - inst->table_td[i - 1] * inst->preemphasis_coefficient;
		}
		inst->table_td[0] = inst->table_td[0] * (1.0f - inst->preemphasis_coefficient);
	}
	hadamard(inst->window, inst->table_td, inst->table_td, inst->window_size);

	if (!inst->raw_energy) {
		signal_log_energy = get_log_energy(inst->table_td, eps, inst->energy_floor, inst->window_size);
	}
	fftr(inst->fft, inst->table_td, inst->table_fd);
	if (inst->use_power) {
		for (int i = 0; i <= inst->padded_window_size / 2; ++i) {
			inst->table_td[i] = inst->table_fd[i].i * inst->table_fd[i].i + inst->table_fd[i].r * inst->table_fd[i].r;
		}
	}
	else {
		for (int i = 0; i < inst->padded_window_size / 2; ++i) {
			inst->table_td[i] = sqrtf(inst->table_fd[i].i * inst->table_fd[i].i + inst->table_fd[i].r * inst->table_fd[i].r);
		}
	}
	float* fbank_start_point;
	if (inst->use_energy && (!inst->htk_compat)) {
		fbank_start_point = fbank_out + 1;
	}
	else {
		fbank_start_point = fbank_out;
	}
	for (int i = 0; i < inst->num_mel_bins; ++i) {
		vec_dot(inst->melbins[i], inst->table_td + inst->bin_start[i], inst->bin_len[i], fbank_start_point + i);
	}
	if (inst->use_log_fbank) {
		for (int i = 0; i < inst->num_mel_bins; ++i) {
			fbank_start_point[i] = fbank_start_point[i] > eps ? fbank_start_point[i] : eps;
			fbank_start_point[i] = logf(fbank_start_point[i]);
		}
	}
	if (inst->use_energy) {
		if (inst->htk_compat) {
			fbank_out[inst->fbank_bins - 1] = signal_log_energy;
		}
		else {
			fbank_out[0] = signal_log_energy;
		}
	}
}

void FbankProcess_block(fbank* inst, float* wav_in, float* out,int16_t block_num) {
	int16_t step = inst->frame_shift * inst->sample_frequency / 1000;
	for (int i = 0; i < block_num; i++) {
		FbankProcess(inst, wav_in + step * i, out + inst->fbank_bins * i);
	}
	
	float* mean = (float*)malloc(sizeof(float) * inst->fbank_bins);
	memset(mean, 0, sizeof(float) * inst->fbank_bins);
	float* p = out;
	if (inst->subtract_mean) {
		for (int i = 0; i < block_num; i++) {
			for (int j = 0; j < inst->fbank_bins; j++) {
				mean[j] += *p;
				p++;
			}
		}
		for (int j = 0; j < inst->fbank_bins; j++) {
			mean[j] /= block_num;
		}
		p = out;
		for (int i = 0; i < block_num; i++) {
			for (int j = 0; j < inst->fbank_bins; j++) {
				*p -= mean[j];
				p++;
			}
		}
	}
	free(mean);
}
void FbankFree(fbank* inst) {
	free(inst->bin_start);
	free(inst->bin_len);
	free(inst->window);
	free(inst->table_td);
	free(inst->fft);
	free(inst->table_fd);
	free(inst->melbins[0]);
	free(inst->melbins);
	free(inst);
}


// 调用示例，实际加入工程的时候下面的代码注释掉， 可以对照下列python语句比对结果
// y = torchaudio.compliance.kaldi.fbank(torch.range(0,1519).unsqueeze(0)/1000.0, num_mel_bins=80,sample_frequency=16000,subtract_mean=True)

void main() {
	float out[80 * 8];
	fbank* obj=FbankInit(80, 16000);
	obj->subtract_mean = 1;
	FbankRegister(obj);
	float testdata[1520];
	for (int i = 0; i < 1520; ++i) {
		testdata[i] = i / 1000.0f;
	}
	FbankProcess_block(obj, testdata, out, 8);
	FbankProcess(obj, testdata, out);
	FbankFree(obj);
}
